from launch import LaunchDescription
from launch.actions import IncludeLaunchDescription, DeclareLaunchArgument, ExecuteProcess, SetEnvironmentVariable
from launch.launch_description_sources import PythonLaunchDescriptionSource
from launch.substitutions import LaunchConfiguration, PathJoinSubstitution, Command
from launch_ros.actions import Node
from launch_ros.substitutions import FindPackageShare
import os
from ament_index_python.packages import get_package_share_directory
import time

def generate_launch_description():
    # 定义可配置参数
    paused = LaunchConfiguration('paused', default='false')
    use_sim_time = LaunchConfiguration('use_sim_time', default='true')
    gui = LaunchConfiguration('gui', default='true')
    headless = LaunchConfiguration('headless', default='false')
    debug = LaunchConfiguration('debug', default='false')
    
    # 获取包路径
    pkg_path = get_package_share_directory('rhea_description')
    
    # 构建自定义世界文件的完整路径
    custom_world_path = os.path.join(pkg_path, 'worlds', 'test_world.world') 
    world_name = LaunchConfiguration('world_name', default=custom_world_path)
    
    # 设置Gazebo模型路径
    gazebo_models_path = os.path.join(pkg_path, 'models')
    gazebo_env = SetEnvironmentVariable(
        name='GAZEBO_MODEL_PATH',
        value=[gazebo_models_path, os.environ.get('GAZEBO_MODEL_PATH', '')]
    )

    # 检查世界文件是否存在
    if not os.path.exists(custom_world_path):
        raise FileNotFoundError(f"世界文件未找到: {custom_world_path}")

    # 启动Gazebo服务器
    gzserver_launch = IncludeLaunchDescription(
        PythonLaunchDescriptionSource([
            PathJoinSubstitution([
                FindPackageShare('gazebo_ros'),
                'launch',
                'gzserver.launch.py'
            ])
        ]),
        launch_arguments={
            'world': world_name,
            'debug': debug,
            'paused': paused,
            'use_sim_time': use_sim_time,
            'headless': headless
        }.items()
    )

    # 启动Gazebo客户端（关键修改）
    gzclient_launch = IncludeLaunchDescription(
        PythonLaunchDescriptionSource([
            PathJoinSubstitution([
                FindPackageShare('gazebo_ros'),
                'launch',
                'gzclient.launch.py'
            ])
        ]),
        launch_arguments={
            'gui': gui
        }.items()
    )

    # 使用Xacro文件
    xacro_file_path = PathJoinSubstitution([
        FindPackageShare('rhea_description'),
        'description',
        'urdf',
        'rhea.urdf.xacro'
    ])
    
    robot_description_content = Command(['xacro ', xacro_file_path])
    
    # 机器人状态发布节点
    robot_state_publisher_node = Node(
        package='robot_state_publisher',
        executable='robot_state_publisher',
        parameters=[{
            'robot_description': robot_description_content,
            'use_sim_time': use_sim_time
        }],
        output='screen'
    )
    
    # 关节状态发布节点
    joint_state_publisher_node = Node(
        package='joint_state_publisher',
        executable='joint_state_publisher',
        parameters=[{'use_sim_time': use_sim_time}],
        output='screen'
    )
    
    # 生成机器人模型
    spawn_entity_node = Node(
        package='gazebo_ros',
        executable='spawn_entity.py',
        arguments=[
            '-entity', 'rhea_' + str(int(time.time())),
            '-topic', 'robot_description',
            '-x', '0', '-y', '0', '-z', '0.5'
        ],
        output='screen'
    )
    
    return LaunchDescription([
        # 环境变量
        SetEnvironmentVariable('GAZEBO_MODEL_PATH', gazebo_models_path),
        SetEnvironmentVariable('LIBGL_ALWAYS_SOFTWARE', '1'),  # 强制软件渲染
        
        # Gazebo服务
        gzserver_launch,
        gzclient_launch,  # 新增的客户端启动
        
        # ROS节点
        robot_state_publisher_node,
        joint_state_publisher_node,
        spawn_entity_node,
        
        # 诊断节点（可选）
        ExecuteProcess(
            cmd=['ros2', 'topic', 'echo', '/joint_states'],
            output='screen',
            shell=True
        )
    ])